package msigraphx.graph;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Vector;

import msigraphx.gibbsSampler.CondDistr;
import msigraphx.gibbsSampler.GibbsSampler;
import msigraphx.gibbsSampler.RandomVariable;
import msigraphx.gibbsSampler.RandomVariableValue;
import msigraphx.gibbsSampler.Sample;
import msigraphx.gibbsSampler.VarDistr;

public class BayesianNetwork {
	private final static int DEFAULT_SAMPLES_NUM = 10000;
	
	private Map<String,NetworkVertex> vertexes;
	private GibbsSampler<String> gibbsSampler;
	
	public BayesianNetwork() {
		vertexes = new HashMap<String,NetworkVertex>();
	}
	
	public void addVertex(NetworkVertex vertex) {
		vertexes.put(vertex.getName(), vertex);
	}
	
	public void removeVertex(String vertexLabel) {
		if(vertexes.containsKey(vertexLabel)) {
			//remove all edges to this vertex and adjust probability distributions
			CondDistr<String> distr;
			for(NetworkVertex vertex : vertexes.get(vertexLabel).getChildren()) {
				distr = vertex.getCondDistr();
				distr.removeDependance(vertexes.get(vertexLabel).getVertexIndex());
			}
			
			//remove vertex
			vertexes.remove(vertexLabel);
		}
	}
	
	public void modifyVertexDistr(String vertexLabel, CondDistr<String> distr) {
		if(vertexes.containsKey(vertexLabel)) {
			NetworkVertex vertex = vertexes.get(vertexLabel);
			vertex.setCondDistr(distr);
		}
	}
	
	public boolean checkConsistency() {
		for(NetworkVertex vertex : vertexes.values()) {
			if(!vertex.checkConsistency()) {
				System.out.println("check failed for vertex: "+vertex.getName()+"  vertexIndex = "+vertex.getVertexIndex());
				return false;
			}
		}
		return true;
	}
	
	private List<Sample<String>> performGibbsSampling(int samplesNum) {
		List<RandomVariable<String>> vertexesList = new Vector<RandomVariable<String>>(vertexes.values().size());
		List<CondDistr<String>> condDistrList = new Vector<CondDistr<String>>(vertexes.values().size());
		
		for(NetworkVertex vertex : vertexes.values()) {
			vertexesList.add(vertex);
			condDistrList.add(vertex.getCondDistr());
		}
		
		for(NetworkVertex vertex : vertexes.values()) {
			vertexesList.set(vertex.getVertexIndex(), vertex);
			condDistrList.set(vertex.getVertexIndex(), vertex.getCondDistr());
		}
		
		gibbsSampler = new GibbsSampler<String>(samplesNum, vertexesList, condDistrList);
		return gibbsSampler.compute();
	}
	
	public VarDistr<String> computeVarDistrByKnownVariables(Map<String,RandomVariableValue<String>> knownVariables, String soughtVariableLabel) {
		int samplesNum = DEFAULT_SAMPLES_NUM;
		return this.computeVarDistrByKnownVariables(knownVariables, soughtVariableLabel, samplesNum);
	}
	
	public VarDistr<String> computeVarDistrByKnownVariables(Map<String,RandomVariableValue<String>> knownVariables, String soughtVariableLabel, Integer samplesNum) {
		List<Sample<String>> samples = performGibbsSampling(samplesNum);
		VarDistr<String> varDistr = new VarDistr<String>();
		NetworkVertex seekVertex = vertexes.get(soughtVariableLabel);
		int tempVertexIndex;
		RandomVariableValue<String> varVal1, varVal2;
		int ommitedNum = 0;
		
		for(Sample<String> sample : samples) {
			boolean ommit = false;
			for(String knownVarLabel : knownVariables.keySet()) {
				varVal2 = knownVariables.get(knownVarLabel);
				tempVertexIndex = vertexes.get(knownVarLabel).getVertexIndex();
				if(!varVal2.equals(sample.values.get(tempVertexIndex))) {
					ommit = true;
					break;
				}
			}
			if(ommit) {
				ommitedNum++;
				continue;
			}
			varVal1 = sample.values.get(seekVertex.getVertexIndex()); //value of sought variable in sample
			if(!varDistr.containsVarValue(varVal1)) {
				varDistr.addValue(varVal1, 1.0f);
			}
			else {
				varDistr.incValue(varVal1);
			}
		}
		varDistr.normalize();
		
		System.out.println("BayesianNetwork.computeVarDistrByKnownVariables()  samplesNum = " + samplesNum + "  samples.size() = " + samples.size()+"  ommitedNum = "+ommitedNum+"  soughtVariableLabel = "+soughtVariableLabel+"  knownVariablesLabes = "+knownVariables.keySet().toString());
		
		return varDistr;
	}

	public NetworkVertex getVertex(String vertexLabel) {
		return vertexes.get(vertexLabel);
	}
	
	public boolean computeDSeparation(List<String> fromVertexesLab, List<String> toVertexesLab, List<String> observVertexesLab) {
		//check args
		Set<String> vertexesKeySet = vertexes.keySet();
		if(!vertexesKeySet.containsAll(fromVertexesLab) 
				|| !vertexesKeySet.containsAll(toVertexesLab) 
				|| !vertexesKeySet.containsAll(observVertexesLab)) {
			throw new IllegalArgumentException();
		}
		
		//init vertexes lists
		Collection<NetworkVertex> fromVertexes = new Vector<NetworkVertex>();
		Collection<NetworkVertex> toVertexes = new Vector<NetworkVertex>();
		Set<NetworkVertex> observVertexes = new HashSet<NetworkVertex>();
		for(String label : fromVertexesLab) {
			fromVertexes.add(vertexes.get(label));
		}
		for(String label : toVertexesLab) {
			toVertexes.add(vertexes.get(label));
		}
		for(String label : observVertexesLab) {
			observVertexes.add(vertexes.get(label));
		}
		
		//get all paths
		Set<List<NetworkVertex>> allPaths = new HashSet<List<NetworkVertex>>();
		List<NetworkVertex> newPath;
		for(NetworkVertex fromVertex : fromVertexes) {
			newPath = new Vector<NetworkVertex>();
			newPath.add(fromVertex);
			allPaths.addAll(this.getAllNewPathsUsingVertex(newPath, toVertexes));
		}
		
		//for each path check if it is blocked
		NetworkVertex prev, actual, next;
		boolean allBlocked = true;
		for(List<NetworkVertex> path : allPaths) {
			if(path.size() <= 2) {
				continue; //always not blocked
			}
			
			boolean pathBlocked = false;
			for(int k=0; k<path.size()-2; ++k) {
				prev = path.get(k);
				actual = path.get(k+1);
				next = path.get(k+2);
				
				if(observVertexes.contains(actual)) {
					if(actual.isParent(prev) || actual.isParent(next)) {
						pathBlocked = true;
					}
				}
				else {
					if(actual.isChild(prev) || actual.isChild(next)) {
						//check all childs
						if(!actual.hasDescendant(observVertexes)) {
							pathBlocked = true;
						}
					}
				}
			}
			
			//System.out.println("pathBlocked = "+pathBlocked);
			if(!pathBlocked) {
				allBlocked = false;
				break;
			}
		}
		
		return allBlocked;
	}
	
	private Set<List<NetworkVertex>> getAllNewPathsUsingVertex(List<NetworkVertex> actualPath, Collection<NetworkVertex> targets) {
		Set<List<NetworkVertex>> paths = new HashSet<List<NetworkVertex>>();
		List<NetworkVertex> newPath;
		
		//String actualLabel = actualPath.get(actualPath.size()-1);
		//NetworkVertex actualVertex = vertexes.get(actualLabel);
		NetworkVertex actualVertex = actualPath.get(actualPath.size()-1);
		
		Set<NetworkVertex> allConnected = actualVertex.getAllConnections();
		for(NetworkVertex connected : allConnected) {
			if(!actualPath.contains(connected)) {
				newPath = new Vector<NetworkVertex>();
				newPath.addAll(actualPath);
				newPath.add(connected);
				paths.addAll(this.getAllNewPathsUsingVertex(newPath, targets));
			}
		}
		
		if(targets.contains(actualVertex)) {
			paths.add(actualPath);
		}
		
		return paths;
	}

	public Map<String, NetworkVertex> getVertexes() {
		return vertexes;
	}
	
	public static BayesianNetwork readFromFile(String fileName) throws IOException {
		return NetworkFromFileReader.readFromFile(fileName);
	}
	
	public static class NetworkFromFileReader {
		private static List<NetworkVertex> vertexes;
		
		public static BayesianNetwork readFromFile(String fileName) throws IOException {
			BayesianNetwork bayesianNetwork = new BayesianNetwork();
			vertexes = new Vector<NetworkVertex>();
			
			FileReader fileReader = new FileReader(fileName);
			BufferedReader reader = new BufferedReader(fileReader);
			
			String vertexesText = "VERTEXES: ";
			String allowedValuesText = "ALLOWED VALUES: ";
			String connectionsText = "CONNECTIONS: ";
			String distrRowText = "DISTR-ROW: ";
			String line;
			while((line=reader.readLine())!=null) {
				if(line.startsWith(vertexesText)) {
					line = line.substring(vertexesText.length());
					readVertexes(line);
				}
				else if(line.startsWith(allowedValuesText)) {
					line = line.substring(allowedValuesText.length()).replaceAll(" ", "");
					readAllowedValues(line);
				}
				else if(line.startsWith(connectionsText)) {
					line = line.substring(connectionsText.length()).replaceAll(" ", "");
					readConnections(line);
				}
				else if(line.startsWith(distrRowText)) {
					line = line.substring(distrRowText.length()).replaceAll(" ", "");
					readDistrRow(line);
				}
				else {
					if(line.charAt(0)!='#') {
						throw new IOException("unsupported format. line = "+line);
					}
				}
			}
			for(NetworkVertex vertex : vertexes) {
				bayesianNetwork.addVertex(vertex);
			}
			
			return bayesianNetwork;
		}

		private static void readDistrRow(String line) {
			List<String> names;
			List<String> names2;
			NetworkVertex actualVertex;
			names = Arrays.asList(line.split(","));
			Float value = Float.parseFloat(names.get(1));
			Integer vertexIndex, allowedValueIndex;
			List<RandomVariableValue<String>> parentsVal = new Vector<RandomVariableValue<String>>();
			if(!"empty".equals(names.get(2).trim().toLowerCase())) {
				List<String> dependantList = names.subList(2, names.size());
				for(String vertexVal : dependantList) {
					names2 = Arrays.asList(vertexVal.split("-"));
					vertexIndex = Integer.parseInt(names2.get(0));
					allowedValueIndex = Integer.parseInt(names2.get(1));
					parentsVal.add(vertexes.get(vertexIndex).getAllowedValue(allowedValueIndex));
				}
			}
			names2 = Arrays.asList(names.get(0).split("-"));
			vertexIndex = Integer.parseInt(names2.get(0));
			allowedValueIndex = Integer.parseInt(names2.get(1));
			actualVertex = vertexes.get(vertexIndex); 
			actualVertex.setCondDistrRow(actualVertex.getAllowedValue(allowedValueIndex), value, parentsVal);
		}

		private static void readConnections(String line) {
			List<String> names;
			List<String> names2;
			names = Arrays.asList(line.split(","));
			for(String connectionDescr : names) {
				names2 = Arrays.asList(connectionDescr.split("-"));
				NetworkVertex vertexFrom = vertexes.get(Integer.parseInt(names2.get(0)));
				NetworkVertex vertexTo = vertexes.get(Integer.parseInt(names2.get(1)));
				vertexFrom.addParent(vertexTo);
			}
		}

		private static void readAllowedValues(String line) {
			List<String> names;
			List<String> names2;
			NetworkVertex actualVertex;
			names = Arrays.asList(line.split("\\|"));
			assert(names.size()==vertexes.size());
			for(int k=0; k<names.size(); ++k) {
				String allowedValues = names.get(k);
				names2 = Arrays.asList(allowedValues.split(","));
				actualVertex = vertexes.get(k);
				for(String valName : names2) {
					actualVertex.addAllowedValue(new RandomVariableValue<String>(actualVertex,valName));
				}
			}
		}

		private static void readVertexes(String line) {
			List<String> names = Arrays.asList(line.split(","));
			for(int k=0; k<names.size(); ++k) {
				vertexes.add(new NetworkVertex(k, names.get(k).trim()));
			}
		}
	}
	
}
